import os
import numpy as np
import tensorflow as tf
from prepocess.utils import get_file, get_batch_tensor, BatchGenerator
from Exp.load_list_to_img_batch import path_list_to_img_batch

from nets.cnn_example import inference, losses, trainning, evaluation
from nets.mobilenet import mobilenet_v2

IMG_W = 224
IMG_H = 224
IMG_CHANNEL = 3
CLASS_NUM = 10
BATCH_SIZE = 32
CAPACITY = 640
MAX_STEP = 5000000
learning_rate = 0.00001
restore_step = 0
ckpt = {}

# 获取批次batch
data_set_dir = './MNIST_lit/'  # 训练样本的读入路径
# train_dir = 'D:/PycharmProjects/qiyi_com/mix_part1'
# train_dir = 'D:/resources/data_set/MNIST_MIX'
logs_train_dir = './train'  # logs存储路径
train, train_label, val, val_label, class_names_to_ids, ids_to_class_names = get_file(data_set_dir)
print(" train files : {} , val files {}".format(len(train), len(val)))

train_set_provider = BatchGenerator(train, train_label, BATCH_SIZE, CLASS_NUM)
val_set_provider = BatchGenerator(val, val_label, BATCH_SIZE, CLASS_NUM)
# train_batch, train_label_batch = get_batch_tensor(train, train_label, IMG_W, IMG_H, BATCH_SIZE, CAPACITY,class_num=CLASS_NUM)
# val_batch, val_label_batch = get_batch_tensor(val, val_label, IMG_W, IMG_H, BATCH_SIZE, CAPACITY,class_num=CLASS_NUM)


graph = tf.Graph()
with graph.as_default():
    with tf.contrib.slim.arg_scope(mobilenet_v2.training_scope()):
        x = tf.placeholder(tf.string, [None, ])
        y_ = tf.placeholder(tf.int32, [None, CLASS_NUM])

        img_batch = path_list_to_img_batch(x, IMG_W, IMG_H, IMG_CHANNEL)
        train_logits, endpoints = mobilenet_v2.mobilenet(img_batch, num_classes=CLASS_NUM, depth_multiplier=1.4)
        train_loss = losses(train_logits, y_)
        train_op = trainning(train_loss, learning_rate)
        acc = evaluation(train_logits, y_, 'acc')
        summary_op = tf.summary.merge_all()

with tf.Session(graph=graph) as sess:
    train_writer = tf.summary.FileWriter(logs_train_dir, sess.graph)
    saver = tf.train.Saver(max_to_keep=3)
    sess.run(tf.global_variables_initializer())
    # 队列监控
    coord = tf.train.Coordinator()  # 设置多线程协调器
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    try:
        for step in np.arange(restore_step, MAX_STEP):
            train_batch = train_set_provider.next_batch()
            if coord.should_stop():
                break
            _, tra_loss, acc_ = sess.run([train_op, train_loss, acc], feed_dict={x: train_batch[0], y_: train_batch[1]})

            if step % 5 == 0:
                val_batch = val_set_provider.next_batch()
                acc_val_ = sess.run(acc, feed_dict={x: val_batch[0], y_: val_batch[1]})
                print('Step %d, train loss = %.2f, train accuracy = %.2f%% val accuracy = %.2f%% ' % (
                    step, tra_loss, acc_ * 100.0, acc_val_ * 100.0))
                # summary_str = sess.run(summary_op)
                # train_writer.add_summary(summary_str, step)
                # checkpoint_path = os.path.join(logs_train_dir, 'mobilenet_v2_1.0_224.ckpt')

    except tf.errors.OutOfRangeError:
        print('Done training -- epoch limit reached')

    finally:
        coord.request_stop()
    coord.join(threads)
    sess.close()
